import json
import time
from threading import Thread
from paramiko import Transport, SSHClient, AutoAddPolicy, RSAKey, DSSKey, ECDSAKey, Ed25519Key


class SSH(object):
    def __init__(self, socket):
        self.message = {'status': 0, 'message': ""}
        self.socket = socket

    def _get_pkey(self, pkey_obj, ssh_key_content, password=None):
        return pkey_obj.from_private_key(ssh_key_content, password=password)

    def shell_to_socket(self):
        while True:
            data = self._channel.recv(4096)
            if not len(data):
                return

            self.message['status'] = 0
            self.message['message'] = data.decode('utf-8')
            message = json.dumps(self.message)
            self.socket.send(message)

    def socket_to_shell(self):
        while True:
            byte_data = self.socket.receive()
            str_data = json.loads(byte_data)
            if str_data["status"] == 0:
                c = str_data["data"]
                if c:
                    self._channel.send(c)
            elif str_data["status"] == 1:
                pty_width = str_data.get("pty_width", 64)
                pty_height = str_data.get("pty_height", 32)
                self._channel.resize_pty(width=pty_width, height=pty_height)

    def connect(self, host, username, password=None, ssh_key_content=None, port=22, timeout=30,
                term='linux', pty_width=80, pty_height=24):
        try:
            transport = Transport(sock=(host, port))
            transport._active_check_timeout = timeout
            if ssh_key_content:
                pkey = self._get_pkey(RSAKey, ssh_key_content, password=password) or \
                       self._get_pkey(DSSKey, ssh_key_content, password=password) or \
                       self._get_pkey(ECDSAKey, ssh_key_content, password=password) or \
                       self._get_pkey(Ed25519Key, ssh_key_content, password=password)

                transport.connect(pkey=pkey)
            else:
                transport.connect(username=username, password=password)

            self._channel = transport.open_session()
            self._channel.get_pty(term=term, width=pty_width, height=pty_height)
            self._channel.invoke_shell()

            for i in range(2):
                recv = self._channel.recv(4096).decode('utf-8')
                message = {
                    "status": 0,
                    "message": recv
                }
                self.socket.send(json.dumps(message))
                time.sleep(0.5)

            Thread(target=self.shell_to_socket).start()

        except Exception:
            self.message['status'] = 2
            self.message['message'] = 'ssh connect faild...'
            message = json.dumps(self.message)
            self.socket.send(message)
            self.socket.close(3001)
